# 导入必要的库
import torch
import torch.nn as nn  # 神经网络模块
import torchvision  # 计算机视觉相关数据集和模型
import torchvision.transforms as transforms  # 数据预处理
from torch.utils.data import DataLoader  # 数据加载器

# 超参数设置
input_size = 784     # 输入层维度（28x28像素展平为784）
hidden_size = 512    # 第一个隐藏层神经元数量
num_classes = 10     # 输出类别数（0-9数字）
num_epochs = 5       # 训练轮次
batch_size = 100     # 每批数据量
learning_rate = 0.001  # 学习率

# MNIST数据集加载
def load_data():
    """
    加载MNIST数据集并进行预处理
    返回:
        train_loader: 训练集数据加载器
        test_loader: 测试集数据加载器
    """
    # 训练集配置
    train_dataset = torchvision.datasets.MNIST(
        root='./data', 
        train=True,
        transform=transforms.Compose([
            transforms.ToTensor(),  # 将PIL图像转换为Tensor[0,1]
            transforms.Normalize((0.1307,), (0.3081,))  # 标准化（均值, 标准差）
        ]),
        download=True  # 自动下载数据集
    )
    
    # 测试集配置（使用相同的归一化参数）
    test_dataset = torchvision.datasets.MNIST(
        root='./data',
        train=False,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
    )
    
    # 创建数据加载器
    return (
        DataLoader(train_dataset, batch_size=batch_size, shuffle=True),  # 训练集打乱顺序
        DataLoader(test_dataset, batch_size=batch_size, shuffle=False)    # 测试集保持顺序
    )

class NeuralNet(nn.Module):
    """三层全连接神经网络模型"""
    def __init__(self, input_size, hidden_size, num_classes):
        """
        初始化网络结构
        参数:
            input_size: 输入维度 
            hidden_size: 隐藏层维度
            num_classes: 输出类别数
        """
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)  # 输入层 → 隐藏层1
        self.relu = nn.ReLU()  # 非线性激活函数
        self.fc2 = nn.Linear(hidden_size, hidden_size//2)  # 隐藏层1 → 隐藏层2
        self.fc3 = nn.Linear(hidden_size//2, num_classes)  # 隐藏层2 → 输出层
        self.dropout = nn.Dropout(0.5)  # 随机失活50%神经元防止过拟合

    def forward(self, x):
        """前向传播流程"""
        x = x.reshape(-1, 28*28)  # 展平图像维度 [batch, 1,28,28] → [batch,784]
        out = self.fc1(x)
        out = self.relu(out)       # 激活函数引入非线性
        out = self.dropout(out)    # 随机失活
        out = self.fc2(out)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc3(out)       # 最终输出（未使用softmax，因CrossEntropyLoss自带）
        return out

def train_model(model, criterion, optimizer, train_loader):
    """
    模型训练函数
    参数:
        model: 要训练的模型
        criterion: 损失函数
        optimizer: 优化器
        train_loader: 训练数据加载器
    """
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
            # 前向传播（自动调用forward方法）
            outputs = model(images)
            loss = criterion(outputs, labels)  # 计算损失
            
            # 反向传播与优化
            optimizer.zero_grad()  # 清空历史梯度
            loss.backward()        # 反向传播计算梯度
            optimizer.step()      # 参数更新
            
            # 每100步打印训练信息
            if (i+1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Loss: {loss.item():.4f}')

def evaluate_model(model, test_loader):
    """
    模型测试
    参数:
        model: 要评估的模型
        test_loader: 测试数据加载器
    """
    with torch.no_grad():  # 禁用梯度计算
        correct = 0
        total = 0
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)  # 获取预测类别
            total += labels.size(0)  # 累计样本总数
            correct += (predicted == labels).sum().item()  # 累计正确数
        
        # 输出整体准确率
        print(f'测试集准确率: {100 * correct / total:.2f}%')

# 主程序入口
if __name__ == '__main__':
    # 数据加载 → 模型初始化 → 训练 → 评估 完整流程
    train_loader, test_loader = load_data()
    # 模型实例化
    model = NeuralNet(input_size, hidden_size, num_classes)
    # 损失函数
    criterion = nn.CrossEntropyLoss() # 交叉熵损失，相当于Softmax+Log+NllLoss
    # 优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 使用Adam优化器
    
    train_model(model, criterion, optimizer, train_loader)
    evaluate_model(model, test_loader)
